import time
from collections import defaultdict
from pathlib import Path
from typing import List, Union

import click
import pandas as pd
from flask.cli import with_appcontext
from pandas.core.frame import DataFrame

from app.models import CtfQuestion, Range, RangeVm
from tools import db
from utils.enums import TextChoices


class IssueColumnField(TextChoices):
    name = "name", "问题名称"
    type = "type", "问题类型"
    score = "score", "问题分数"
    answer = "answer", "问题答案"
    describe = "describe", "问题描述"
    belong = "belong", "问题归属"
    filepath = "filepath", "问题附件路径"


def add_system_issue(df: DataFrame):
    """添加系统赛题"""
    try:
        fields = IssueColumnField.values
        with db.auto_commit():
            issue_list = [CtfQuestion(**dict(zip(fields, item))) for item in df.values]
            db.session.bulk_save_objects(issue_list)
    except Exception as e:
        print("导入系统赛题失败: ", e)
        db.session.rollback()
        db.session.close()
        raise e
    else:
        print(f"已导入系统内置公开赛题：{len(issue_list)} 道")


def bind_vm_and_issue(df: DataFrame):
    """绑定系统内置靶场中虚拟机内置赛题"""

    # 按靶场-虚拟机-赛题构建数据
    fields = IssueColumnField.values
    vm_bind_issue_info = defaultdict(lambda: defaultdict(list))
    for item in df.values:
        vm_bind_issue_info[item[0]][item[1]].append(dict(zip(fields, item[2:])))

    # 创建并与对应内置虚拟机建立绑定关系
    for range_type, vm_bind_info in vm_bind_issue_info.items():
        # 初始化系统时，此处应只能查询到一个场景数据
        _range: Range = Range.query.filter(
            Range.range_template_type == range_type, Range.is_system == 1
        ).first()
        print(f"检索{_range.name}靶场...")
        # 获取到该场景中所有的vm
        vm_list: List[RangeVm] = _range.rangeVmList[0].vmList
        for vm in vm_list:
            print(f"\t对虚拟机{vm.vm_name}进行操作")
            # 获取当前vm_name对应的赛题
            _vm_bind: Union[List, None] = vm_bind_info.get(vm.vm_name, None)
            if not _vm_bind:
                continue
            # 导入虚拟机内置赛题并建立关系
            issue_list = [CtfQuestion(**item, bind_vm=[vm]) for item in _vm_bind]
            db.session.add_all(issue_list)
            db.session.flush()
        print(f"{_range.name}靶场操作完成")


@click.command(
    "add_issue",
    context_settings={"ignore_unknown_options": True},
    short_help="添加系统内置赛题, 并对虚拟机内置赛题建立绑定关系",
)
@click.option(
    "--bind",
    type=bool,
    default=True,
    help="对于虚拟机内置赛题，自动建立绑定关系",
)
@with_appcontext
def add_issue(bind=True):
    """添加系统内置靶场"""

    print("该命令应在 `flask add_range` 命令之后执行")
    time.sleep(2)

    system_issue = Path(__file__).parent / "lib/system_issue.xlsx"

    # 导入系统内置赛题
    if system_issue.exists():
        df: DataFrame = pd.read_excel(system_issue, keep_default_na=False)
        add_system_issue(df)

    # 判断是否需要导入虚拟机内置赛题
    if bind is False:
        return

    # 导入系统虚拟机内置赛题并建立与对应虚拟机绑定关系
    system_bind_vm_issue = Path(__file__).parent / "lib/system_bind_vm_issue.xlsx"
    if system_bind_vm_issue.exists():
        try:
            df: DataFrame = pd.read_excel(system_bind_vm_issue, keep_default_na=False)
            bind_vm_and_issue(df)
            db.session.commit()
        except Exception as e:
            print("建立赛题与内置虚拟机绑定关系失败:", e)
            db.session.rollback()
            db.session.close()
